{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this tutorial, we will show four ways to use generator:\n",
    "1. Use directly by setting up a ``ModelClient``.\n",
    "2. Customize prompt template using ``jinjia2``.\n",
    "3. Try different models.\n",
    "4. Use ``acall`` to do mutiple asynchronous calls for speed up."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In default, the generator uses a default prompt template. It has these varaibles:\n",
    "\n",
    "LIGHTRAG_DEFAULT_PROMPT_ARGS = [\n",
    "    \"task_desc_str\",\n",
    "    \"output_format_str\",\n",
    "    \"tools_str\",\n",
    "    \"examples_str\",\n",
    "    \"chat_history_str\",\n",
    "    \"context_str\",\n",
    "    \"steps_str\",\n",
    "    \"input_str\",\n",
    "    \"output_str\",\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# first, let's set up the library log just in case, in default at INFO level\n",
    "from utils.logger import enable_library_logging\n",
    "enable_library_logging()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2024-06-09 22:06:49 - INFO - [prompt_builder.py:82:__init__] - Prompt has variables: ['chat_history_str', 'task_desc_str', 'steps_str', 'examples_str', 'tools_str', 'context_str', 'output_str', 'output_format_str', 'input_str']\n",
      "2024-06-09 22:06:49 - INFO - [generator.py:194:call] - prompt_kwargs: {'input_str': 'What is the capital of France?'}\n",
      "2024-06-09 22:06:49 - INFO - [generator.py:195:call] - model_kwargs: {}\n",
      "2024-06-09 22:06:49 - INFO - [openai_client.py:122:call] - api_kwargs: {'model': 'gpt-3.5-turbo', 'messages': [{'role': 'system', 'content': '<Inputs>\\nWhat is the capital of France?\\n</Inputs>'}]}\n",
      "2024-06-09 22:06:49 - INFO - [_client.py:1026:_send_single_request] - HTTP Request: POST https://api.openai.com/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
      "2024-06-09 22:06:49 - INFO - [generator.py:203:call] - output: GeneratorOutput(data='The capital of France is Paris.', error=None, raw_response='The capital of France is Paris.')\n",
      "GeneratorOutput(data='The capital of France is Paris.', error=None, raw_response='The capital of France is Paris.')\n"
     ]
    }
   ],
   "source": [
    "from lightrag.core import Generator\n",
    "from lightrag.components.model_client import OpenAIClient\n",
    "from lightrag.utils import setup_env # ensure you have .env with OPENAI_API_KEY\n",
    "\n",
    "query = \"What is the capital of France?\"\n",
    "model_kwargs = {\n",
    "    \"model\": \"gpt-3.5-turbo\"\n",
    "}\n",
    "generator = Generator(model_client=OpenAIClient(), model_kwargs=model_kwargs)\n",
    "prompt_kwargs = {\n",
    "    \"input_str\": query,\n",
    "}\n",
    "# run the generator\n",
    "output = generator(prompt_kwargs=prompt_kwargs)\n",
    "print(output)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The logging clearly shows us what we sent to OpenAI."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Prompt:\n",
      "\n",
      "<Inputs>\n",
      "What is the capital of France?\n",
      "</Inputs>\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# lets see the prompt, it is quite minimal\n",
    "generator.print_prompt(**prompt_kwargs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Writing your template is easy. Let us use our own template. Let's say, we want to set up our AI with a sense of humor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2024-06-09 22:09:43 - INFO - [prompt_builder.py:82:__init__] - Prompt has variables: ['input_str']\n",
      "2024-06-09 22:09:43 - INFO - [generator.py:194:call] - prompt_kwargs: {'input_str': 'What is the capital of France?'}\n",
      "2024-06-09 22:09:43 - INFO - [generator.py:195:call] - model_kwargs: {}\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2024-06-09 22:09:43 - INFO - [openai_client.py:122:call] - api_kwargs: {'model': 'gpt-3.5-turbo', 'messages': [{'role': 'system', 'content': '<SYS> Your are an assistant with a great sense of humor.</SYS> User: What is the capital of France?. You:'}]}\n",
      "2024-06-09 22:09:44 - INFO - [_client.py:1026:_send_single_request] - HTTP Request: POST https://api.openai.com/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
      "2024-06-09 22:09:44 - INFO - [generator.py:203:call] - output: GeneratorOutput(data='The capital of France is Paris. It\\'s the city of love, pastries, and baguettes, so pack your beret and get ready to say \"ooh la la!\"', error=None, raw_response='The capital of France is Paris. It\\'s the city of love, pastries, and baguettes, so pack your beret and get ready to say \"ooh la la!\"')\n",
      "GeneratorOutput(data='The capital of France is Paris. It\\'s the city of love, pastries, and baguettes, so pack your beret and get ready to say \"ooh la la!\"', error=None, raw_response='The capital of France is Paris. It\\'s the city of love, pastries, and baguettes, so pack your beret and get ready to say \"ooh la la!\"')\n"
     ]
    }
   ],
   "source": [
    "template = \"\"\"<SYS> Your are an assistant with a great sense of humor.</SYS> User: {{input_str}}. You:\"\"\"\n",
    "\n",
    "generator2 = Generator(model_client=OpenAIClient(), model_kwargs=model_kwargs, template=template)\n",
    "response = generator2(prompt_kwargs=prompt_kwargs)\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2024-06-09 22:12:28 - INFO - [prompt_builder.py:82:__init__] - Prompt has variables: ['input_str']\n",
      "2024-06-09 22:12:28 - INFO - [generator.py:194:call] - prompt_kwargs: {'input_str': 'What is the capital of France?'}\n",
      "2024-06-09 22:12:28 - INFO - [generator.py:195:call] - model_kwargs: {}\n",
      "2024-06-09 22:12:29 - INFO - [_client.py:1026:_send_single_request] - HTTP Request: POST https://api.groq.com/openai/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
      "2024-06-09 22:12:29 - INFO - [generator.py:203:call] - output: GeneratorOutput(data='Bonjour! The capital of France is indeed Paris! But did you know that Paris is also the city of love, famous for its Eiffel Tower, Louvre Museum, and, of course, croissants...', error=None, raw_response='Bonjour! The capital of France is indeed Paris! But did you know that Paris is also the city of love, famous for its Eiffel Tower, Louvre Museum, and, of course, croissants...')\n",
      "GeneratorOutput(data='Bonjour! The capital of France is indeed Paris! But did you know that Paris is also the city of love, famous for its Eiffel Tower, Louvre Museum, and, of course, croissants...', error=None, raw_response='Bonjour! The capital of France is indeed Paris! But did you know that Paris is also the city of love, famous for its Eiffel Tower, Louvre Museum, and, of course, croissants...')\n"
     ]
    }
   ],
   "source": [
    "# Let us use llama3 from groq\n",
    "from lightrag.components.model_client import GroqAPIClient\n",
    "\n",
    "groq_model_kwargs = {\"model\": \"llama3-8b-8192\"}\n",
    "generator3 = Generator(model_client=GroqAPIClient(), model_kwargs=groq_model_kwargs, template=template)\n",
    "\n",
    "response = generator3(prompt_kwargs=prompt_kwargs)\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2024-06-09 22:20:08 - INFO - [generator.py:217:acall] - prompt_kwargs: {'input_str': 'What is the capital of France?'}\n",
      "2024-06-09 22:20:08 - INFO - [generator.py:218:acall] - model_kwargs: {}\n",
      "2024-06-09 22:20:08 - INFO - [generator.py:217:acall] - prompt_kwargs: {'input_str': 'What is the capital of France?'}\n",
      "2024-06-09 22:20:08 - INFO - [generator.py:218:acall] - model_kwargs: {}\n",
      "2024-06-09 22:20:08 - INFO - [generator.py:217:acall] - prompt_kwargs: {'input_str': 'What is the capital of France?'}\n",
      "2024-06-09 22:20:08 - INFO - [generator.py:218:acall] - model_kwargs: {}\n",
      "2024-06-09 22:20:08 - INFO - [generator.py:217:acall] - prompt_kwargs: {'input_str': 'What is the capital of France?'}\n",
      "2024-06-09 22:20:08 - INFO - [generator.py:218:acall] - model_kwargs: {}\n",
      "2024-06-09 22:20:08 - INFO - [generator.py:217:acall] - prompt_kwargs: {'input_str': 'What is the capital of France?'}\n",
      "2024-06-09 22:20:08 - INFO - [generator.py:218:acall] - model_kwargs: {}\n",
      "2024-06-09 22:20:08 - INFO - [generator.py:217:acall] - prompt_kwargs: {'input_str': 'What is the capital of France?'}\n",
      "2024-06-09 22:20:08 - INFO - [generator.py:218:acall] - model_kwargs: {}\n",
      "2024-06-09 22:20:08 - INFO - [generator.py:217:acall] - prompt_kwargs: {'input_str': 'What is the capital of France?'}\n",
      "2024-06-09 22:20:08 - INFO - [generator.py:218:acall] - model_kwargs: {}\n",
      "2024-06-09 22:20:08 - INFO - [generator.py:217:acall] - prompt_kwargs: {'input_str': 'What is the capital of France?'}\n",
      "2024-06-09 22:20:08 - INFO - [generator.py:218:acall] - model_kwargs: {}\n",
      "2024-06-09 22:20:08 - INFO - [generator.py:217:acall] - prompt_kwargs: {'input_str': 'What is the capital of France?'}\n",
      "2024-06-09 22:20:08 - INFO - [generator.py:218:acall] - model_kwargs: {}\n",
      "2024-06-09 22:20:08 - INFO - [generator.py:217:acall] - prompt_kwargs: {'input_str': 'What is the capital of France?'}\n",
      "2024-06-09 22:20:08 - INFO - [generator.py:218:acall] - model_kwargs: {}\n",
      "2024-06-09 22:20:09 - INFO - [_client.py:1773:_send_single_request] - HTTP Request: POST https://api.groq.com/openai/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
      "2024-06-09 22:20:09 - INFO - [generator.py:225:acall] - output: GeneratorOutput(data=\"Bonjour! The capital of France is Paris, of course! But did you know that Paris is actually the City of Love? At least, that's what the tourists keep telling me.\", error=None, raw_response=\"Bonjour! The capital of France is Paris, of course! But did you know that Paris is actually the City of Love? At least, that's what the tourists keep telling me.\")\n",
      "2024-06-09 22:20:09 - INFO - [_client.py:1773:_send_single_request] - HTTP Request: POST https://api.groq.com/openai/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
      "2024-06-09 22:20:09 - INFO - [_client.py:1773:_send_single_request] - HTTP Request: POST https://api.groq.com/openai/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
      "2024-06-09 22:20:09 - INFO - [_client.py:1773:_send_single_request] - HTTP Request: POST https://api.groq.com/openai/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
      "2024-06-09 22:20:09 - INFO - [_client.py:1773:_send_single_request] - HTTP Request: POST https://api.groq.com/openai/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
      "2024-06-09 22:20:09 - INFO - [_client.py:1773:_send_single_request] - HTTP Request: POST https://api.groq.com/openai/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
      "2024-06-09 22:20:09 - INFO - [generator.py:225:acall] - output: GeneratorOutput(data='Bien sûr! The capital of France is... (drumroll, please)... Paris!', error=None, raw_response='Bien sûr! The capital of France is... (drumroll, please)... Paris!')\n",
      "2024-06-09 22:20:09 - INFO - [generator.py:225:acall] - output: GeneratorOutput(data='Bonjour!\\n\\nThe capital of France is none other than Paris! (or as the French like to call it, \"The City of Love and Croissants\").\\n\\nYou know what they say: Paris is always a good idea...', error=None, raw_response='Bonjour!\\n\\nThe capital of France is none other than Paris! (or as the French like to call it, \"The City of Love and Croissants\").\\n\\nYou know what they say: Paris is always a good idea...')\n",
      "2024-06-09 22:20:09 - INFO - [generator.py:225:acall] - output: GeneratorOutput(data=\"Bonjour! The capital of France is Paris, of course! But did you know that it's also the City of Love, where the Eiffel Tower stands tall and the croissants are always fresh?\", error=None, raw_response=\"Bonjour! The capital of France is Paris, of course! But did you know that it's also the City of Love, where the Eiffel Tower stands tall and the croissants are always fresh?\")\n",
      "2024-06-09 22:20:09 - INFO - [generator.py:225:acall] - output: GeneratorOutput(data='Bonjour! The capital of France is Paris, naturally!', error=None, raw_response='Bonjour! The capital of France is Paris, naturally!')\n",
      "2024-06-09 22:20:09 - INFO - [generator.py:225:acall] - output: GeneratorOutput(data='Bonjour! The capital of France is... (drumroll please)... Paris! And if you\\'re feeling fancy, you can also say \"Je suis à Paris, comment allez-vous?\" which roughly translates to \"I\\'m in Paris, how are you?\" but actually means \"I\\'m in Paris, drop everything and let\\'s go party!\"', error=None, raw_response='Bonjour! The capital of France is... (drumroll please)... Paris! And if you\\'re feeling fancy, you can also say \"Je suis à Paris, comment allez-vous?\" which roughly translates to \"I\\'m in Paris, how are you?\" but actually means \"I\\'m in Paris, drop everything and let\\'s go party!\"')\n",
      "2024-06-09 22:20:09 - INFO - [_client.py:1773:_send_single_request] - HTTP Request: POST https://api.groq.com/openai/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
      "2024-06-09 22:20:09 - INFO - [generator.py:225:acall] - output: GeneratorOutput(data=\"Bonjour! The capital of France is... (drumroll please)... PARIS! Nope, just kidding, that would be too easy! Seriously though, the capital of France is indeed Paris! But let's keep the suspense going, shall we?\", error=None, raw_response=\"Bonjour! The capital of France is... (drumroll please)... PARIS! Nope, just kidding, that would be too easy! Seriously though, the capital of France is indeed Paris! But let's keep the suspense going, shall we?\")\n",
      "2024-06-09 22:20:09 - INFO - [_client.py:1773:_send_single_request] - HTTP Request: POST https://api.groq.com/openai/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
      "2024-06-09 22:20:09 - INFO - [generator.py:225:acall] - output: GeneratorOutput(data=\"Bonjour! The capital of France is... (drumroll please)... Paris! But seriously, you can't even imagine how much cheese goes into making decisions in that city.\", error=None, raw_response=\"Bonjour! The capital of France is... (drumroll please)... Paris! But seriously, you can't even imagine how much cheese goes into making decisions in that city.\")\n",
      "2024-06-09 22:20:09 - INFO - [_client.py:1773:_send_single_request] - HTTP Request: POST https://api.groq.com/openai/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
      "2024-06-09 22:20:09 - INFO - [generator.py:225:acall] - output: GeneratorOutput(data='Bonjour! The capital of France is Paris, bien sûr! But did you know that Paris is often referred to as the \"City of Love\"? Maybe it\\'s because the Eiffel Tower is the most romantic landmark in the world... or maybe it\\'s because the French have a certain... je ne sais quoi when it comes to romance.', error=None, raw_response='Bonjour! The capital of France is Paris, bien sûr! But did you know that Paris is often referred to as the \"City of Love\"? Maybe it\\'s because the Eiffel Tower is the most romantic landmark in the world... or maybe it\\'s because the French have a certain... je ne sais quoi when it comes to romance.')\n",
      "2024-06-09 22:20:09 - INFO - [_client.py:1773:_send_single_request] - HTTP Request: POST https://api.groq.com/openai/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
      "2024-06-09 22:20:09 - INFO - [generator.py:225:acall] - output: GeneratorOutput(data='Bonjour! The capital of France is Paris, darling! But did you know that Paris is also the city of love, art, and croissants?', error=None, raw_response='Bonjour! The capital of France is Paris, darling! But did you know that Paris is also the city of love, art, and croissants?')\n",
      "Time taken for 10 async calls: 0.8704450130462646\n",
      "[GeneratorOutput(data='Bonjour! The capital of France is Paris, darling! But did you know that Paris is also the city of love, art, and croissants?', error=None, raw_response='Bonjour! The capital of France is Paris, darling! But did you know that Paris is also the city of love, art, and croissants?'), GeneratorOutput(data=\"Bonjour! The capital of France is... (drumroll please)... PARIS! Nope, just kidding, that would be too easy! Seriously though, the capital of France is indeed Paris! But let's keep the suspense going, shall we?\", error=None, raw_response=\"Bonjour! The capital of France is... (drumroll please)... PARIS! Nope, just kidding, that would be too easy! Seriously though, the capital of France is indeed Paris! But let's keep the suspense going, shall we?\"), GeneratorOutput(data='Bien sûr! The capital of France is... (drumroll, please)... Paris!', error=None, raw_response='Bien sûr! The capital of France is... (drumroll, please)... Paris!'), GeneratorOutput(data=\"Bonjour! The capital of France is Paris, of course! But did you know that Paris is actually the City of Love? At least, that's what the tourists keep telling me.\", error=None, raw_response=\"Bonjour! The capital of France is Paris, of course! But did you know that Paris is actually the City of Love? At least, that's what the tourists keep telling me.\"), GeneratorOutput(data=\"Bonjour! The capital of France is... (drumroll please)... Paris! But seriously, you can't even imagine how much cheese goes into making decisions in that city.\", error=None, raw_response=\"Bonjour! The capital of France is... (drumroll please)... Paris! But seriously, you can't even imagine how much cheese goes into making decisions in that city.\"), GeneratorOutput(data='Bonjour! The capital of France is Paris, naturally!', error=None, raw_response='Bonjour! The capital of France is Paris, naturally!'), GeneratorOutput(data=\"Bonjour! The capital of France is Paris, of course! But did you know that it's also the City of Love, where the Eiffel Tower stands tall and the croissants are always fresh?\", error=None, raw_response=\"Bonjour! The capital of France is Paris, of course! But did you know that it's also the City of Love, where the Eiffel Tower stands tall and the croissants are always fresh?\"), GeneratorOutput(data='Bonjour! The capital of France is... (drumroll please)... Paris! And if you\\'re feeling fancy, you can also say \"Je suis à Paris, comment allez-vous?\" which roughly translates to \"I\\'m in Paris, how are you?\" but actually means \"I\\'m in Paris, drop everything and let\\'s go party!\"', error=None, raw_response='Bonjour! The capital of France is... (drumroll please)... Paris! And if you\\'re feeling fancy, you can also say \"Je suis à Paris, comment allez-vous?\" which roughly translates to \"I\\'m in Paris, how are you?\" but actually means \"I\\'m in Paris, drop everything and let\\'s go party!\"'), GeneratorOutput(data='Bonjour! The capital of France is Paris, bien sûr! But did you know that Paris is often referred to as the \"City of Love\"? Maybe it\\'s because the Eiffel Tower is the most romantic landmark in the world... or maybe it\\'s because the French have a certain... je ne sais quoi when it comes to romance.', error=None, raw_response='Bonjour! The capital of France is Paris, bien sûr! But did you know that Paris is often referred to as the \"City of Love\"? Maybe it\\'s because the Eiffel Tower is the most romantic landmark in the world... or maybe it\\'s because the French have a certain... je ne sais quoi when it comes to romance.'), GeneratorOutput(data='Bonjour!\\n\\nThe capital of France is none other than Paris! (or as the French like to call it, \"The City of Love and Croissants\").\\n\\nYou know what they say: Paris is always a good idea...', error=None, raw_response='Bonjour!\\n\\nThe capital of France is none other than Paris! (or as the French like to call it, \"The City of Love and Croissants\").\\n\\nYou know what they say: Paris is always a good idea...')]\n"
     ]
    }
   ],
   "source": [
    "# Lets do 10 async calls at once, lets use GroqAPIClient\n",
    "import nest_asyncio # import asyncio, use nest_asyncio.apply() if you are in jupyter notebook\n",
    "import asyncio\n",
    "nest_asyncio.apply()\n",
    "\n",
    "import time\n",
    "from typing import List\n",
    "\n",
    "async def make_async_calls(queries: List[str]):\n",
    "    calls = [generator3.acall(prompt_kwargs={\"input_str\": query}) for query in queries]\n",
    "    responses = await asyncio.gather(*calls)\n",
    "    return responses\n",
    "\n",
    "queries = [query] * 10\n",
    "start = time.time()\n",
    "responses = asyncio.run(make_async_calls(queries))\n",
    "print(f\"Time taken for 10 async calls: {time.time() - start}\")\n",
    "print(responses)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "my-project-kernel",
   "language": "python",
   "name": "my-project-kernel"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
